/** @file dataset_parser_test.cc
 *
 *  @author Dongryeol Lee (drselee@gmail.com)
 */

#include <armadillo>
#include <fstream>

#include "hdf5.h"
#include "hdf5_hl.h"

#include "core/data/csv_parser.h"
#include "core/data/hdf5_dataset.h"
#include "core/data/matrix_dataset.h"
#include "core/data/sparse_dataset.h"
#include "core/data/sqlite_dataset.h"
#include "core/math/math.h"
#include "core/util/random.h"

namespace core {
namespace data {
class DatasetParserTest {
  private:
    static void InsertRandomSpaces_(
      std::ofstream &foutput, core::util::Random &random) {
      int num_random_spaces = random.Integer(5);
      for(int k = 0; k < num_random_spaces; k++) {
        foutput << " ";
      }
    }

  public:

    static void Hdf5Test() {
      std::cerr << "Hdf5Test: Testing the single process parsing...\n";

      core::util::Random random;
      hid_t       file_id;
      std::vector< hsize_t > dims(2, 0);
      dims[0] = random.Integer(5, 10);
      dims[1] = random.Integer(5, 10);
      std::vector<double> data(dims[0] * dims[1], 0.0);
      for(unsigned int i = 0; i < data.size(); i++) {
        data[i] = i;
      }
      herr_t      status;

      // create a HDF5 file
      file_id = H5Fcreate(
                  "ex_lite1.h5", H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);

      // create and write an integer type dataset named "dset"
      status = H5LTmake_dataset(
                 file_id, "/dset", 2, dims.data(),
                 H5T_NATIVE_DOUBLE, data.data());

      // close file
      status = H5Fclose(file_id);

      // Re-read.
      core::data::Hdf5Dataset data_reread;
      data_reread.Init("ex_lite1.h5/dset");
    }

    static void SparseTest() {

      std::cerr << "SparseTest: Testing the single process parsing...\n";
      core::util::Random random;

      int num_attributes = random.Integer(5, 10);
      int num_points = random.Integer(80, 120);
      std::string test_sparse_file_name("sparse_test.csv");
      std::ofstream foutput(
        test_sparse_file_name.c_str(),
        std::ios_base::out | std::ios_base::trunc);

      arma::mat random_mat;
      random_mat.zeros(num_attributes, num_points);
      for(int j = 0; j < num_points; j++) {
        if(random.Integer(2) < 1) {
          foutput << "-1 ";
          random_mat.at(num_attributes - 1, j) = -1;
        }
        else {
          foutput << "+1 ";
          random_mat.at(num_attributes - 1, j) = 1;
        }
        InsertRandomSpaces_(foutput, random);
        for(int i = 0; i < num_attributes - 1; i++) {
          if(random.Integer(2) < 1) {
            InsertRandomSpaces_(foutput, random);
            double random_num = random.Uniform();
            random_mat.at(i, j) = random_num;
            int index = i + 1;
            foutput << index << ":" << random_num << " ";
            InsertRandomSpaces_(foutput, random);
          }
        }
        foutput << "\n";
      }
      foutput.close();

      // Try to parse the randomly generated sparse data file.
      core::data::SparseDataset dataset;
      if((! dataset.Init(test_sparse_file_name)) ||
          dataset.num_attributes() != num_attributes) {
        std::cerr << "Failed to be parsed...\n";
        printf("Hi %d %d\n", dataset.num_attributes(), num_attributes);
        exit(0);
      }

      arma::mat dense_matrix_rep;
      dataset.as_dense_matrix(&dense_matrix_rep);
      for(int j = 0; j < dataset.num_points(); j++) {
        for(int i = 0; i < dataset.num_attributes(); i++) {
          if(! core::math::ApproxEqual(
                dense_matrix_rep.at(i, j) , random_mat.at(i, j))) {
            std::cerr << "Failed at (" << i << ", " << j << ")...\n";
            printf("%g vs %g\n", dense_matrix_rep.at(i, j),
                   random_mat.at(i, j));
            exit(0);
          }
        }
      }
      std::cerr << "SparseTest: all tests passed!\n";
    }

    static void CSVTest() {

      std::cerr << "Test: Testing the single process parsing...\n";
      core::util::Random random;

      int num_attributes = random.Integer(2, 8);
      int num_points = random.Integer(50, 100);
      std::string test_csv_file_name("test.csv");
      std::ofstream foutput(
        test_csv_file_name.c_str(), std::ios_base::out | std::ios_base::trunc);
      for(int i = 0; i < num_attributes; i++) {
        InsertRandomSpaces_(foutput, random);
        foutput << "att_" << i;
        InsertRandomSpaces_(foutput, random);
        if(i < num_attributes - 1) {
          foutput << ",";
        }
      }
      foutput << "\n";
      arma::mat random_mat;
      random_mat.zeros(num_attributes, num_points);
      for(int j = 0; j < num_points; j++) {
        for(int i = 0; i < num_attributes; i++) {
          InsertRandomSpaces_(foutput, random);
          double random_num = random.Uniform();
          random_mat.at(i, j) = random_num;
          foutput << random_num;
          InsertRandomSpaces_(foutput, random);
          if(i < num_attributes - 1) {
            foutput << ",";
          }
        }
        foutput << "\n";
      }
      foutput.close();

      // Try to parse the randomly generated CSV file.
      core::data::MatrixDataset dataset;
      if((! dataset.Init(test_csv_file_name)) ||
          dataset.num_attributes() != num_attributes) {
        std::cerr << "Failed to be parsed...\n";
        exit(0);
      }

      for(int j = 0; j < num_points; j++) {
        for(int i = 0; i < num_attributes; i++) {
          if(! core::math::ApproxEqual(
                random_mat.at(i, j), dataset.get_instance(j)[i])) {
            std::cerr << "Not equal...\n";
            exit(0);
          }
        }
      }
      std::cerr << "CSVTest: all tests passed!\n\n";
    }

    static void ParallelCSVTest(boost::mpi::communicator &world) {

      if(world.rank() == 0) {
        std::cerr << "ParallelCSVTest: testing parallel reading with " <<
                  world.size() << " processes...\n";
      }

      // Paralel reading.
      std::string test_file_name("test.csv");
      core::data::MatrixDataset data;
      data.Init(world, test_file_name, true);

      // Now master collects all data back and checks manually.
      if(world.rank() == 0) {
        core::data::MatrixDataset whole_data;
        whole_data.Init(test_file_name);
        arma::mat *subdata_list = new arma::mat[ world.size() ];
        boost::mpi::gather(world, data.matrix_, subdata_list, 0);
        int offset_num_points = 0;
        for(int k = 0; k < world.size(); k++) {
          for(int j = 0; j < static_cast<int>(subdata_list[k].n_cols); j++) {
            for(int i = 0; i < static_cast<int>(subdata_list[k].n_rows); i++) {
              if(! core::math::ApproxEqual(
                    subdata_list[k].at(i, j),
                    whole_data.matrix_.at(i, offset_num_points + j))) {
                std::cerr << "Bug in parallel parsing at " << k << " " <<
                          i << " " << j << "...\n";
                exit(0);
              }
            }
          }
          offset_num_points += static_cast<int>(subdata_list[k].n_cols);
        }
        delete[] subdata_list;
      }
      else {
        boost::mpi::gather(world, data.matrix_, 0);
      }

      if(world.rank() == 0) {
        std::cerr << "ParallelCSVTest: all tests passed!\n\n";
      }
    }

    static void ParallelSparseTest(boost::mpi::communicator &world) {

      if(world.rank() == 0) {
        std::cerr << "ParallelSparseTest: testing parallel reading with " <<
                  world.size() << " processes...\n";
      }

      // Paralel reading.
      std::string test_file_name("sparse_test.csv");
      core::data::SparseDataset data;
      data.Init(world, test_file_name, true);

      // Now master collects all data back and checks manually.
      if(world.rank() == 0) {
        core::data::SparseDataset whole_data;
        arma::mat whole_data_dense_matrix;
        whole_data.Init(test_file_name);
        whole_data.as_dense_matrix(& whole_data_dense_matrix);
        arma::mat *subdata_list = new arma::mat[ world.size() ];
        arma::mat self_dense_matrix;
        data.as_dense_matrix(& self_dense_matrix);
        boost::mpi::gather(world, self_dense_matrix, subdata_list, 0);
        int offset_num_points = 0;
        for(int k = 0; k < world.size(); k++) {
          for(int j = 0; j < static_cast<int>(subdata_list[k].n_cols); j++) {
            for(int i = 0; i < static_cast<int>(subdata_list[k].n_rows); i++) {
              if(! core::math::ApproxEqual(
                    subdata_list[k].at(i, j),
                    whole_data_dense_matrix.at(i, offset_num_points + j))) {
                std::cerr << "Bug in parallel parsing at " << k << " " <<
                          i << " " << j << "...\n";
                exit(0);
              }
            }
          }
          offset_num_points += static_cast<int>(subdata_list[k].n_cols);
        }
        delete[] subdata_list;
      }
      else {
        arma::mat self_dense_matrix;
        data.as_dense_matrix(& self_dense_matrix);
        boost::mpi::gather(world, self_dense_matrix, 0);
      }

      if(world.rank() == 0) {
        std::cerr << "ParallelSparseTest: all tests passed!\n\n";
      }
    }

    static void SqliteDatasetTest() {

      // Delete and re-create the test dataset.
      const char *file_name = "test.db";
      sqlite3 *db_connection = NULL;
      std::ifstream ifile(file_name);
      if(ifile) {
        std::remove(file_name);
      }
      sqlite3_open(file_name, & db_connection);
      sqlite3_exec(
        db_connection,
        "CREATE TABLE IF NOT EXISTS t1(uname TEXT PRIMARY KEY,pass TEXT NOT NULL,activated INTEGER)", 0, 0, 0);
      sqlite3_exec(
        db_connection, "INSERT INTO t1 VALUES('user1','password1',1)", 0, 0, 0);
      sqlite3_exec(
        db_connection, "INSERT INTO t1 VALUES('user2','password2',0)", 0, 0, 0);
      sqlite3_exec(
        db_connection, "INSERT INTO t1 VALUES('user3','password3',0)", 0, 0, 0);
      sqlite3_close(db_connection);

      core::data::SqliteDataset dataset;
      dataset.Init(std::string("test.db:t1"));
    }
};
}
}

int main(int argc, char *argv[]) {

  // Initialize boost MPI.
  boost::mpi::environment env(argc, argv);
  boost::mpi::communicator world;
  if(world.rank() == 0) {
    core::data::DatasetParserTest::CSVTest();
    core::data::DatasetParserTest::SparseTest();
    core::data::DatasetParserTest::Hdf5Test();
  }
  world.barrier();

  // Test the parallel file reading.
  core::data::DatasetParserTest::ParallelCSVTest(world);
  core::data::DatasetParserTest::ParallelSparseTest(world);

  // Test the Sqlite reading.
  if(world.rank() == 0) {
    core::data::DatasetParserTest::SqliteDatasetTest();
  }
  return 0;
}
